{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Flower Quickstart (Simulation with TensorFlow/Keras)\n",
    "\n",
    "Welcome to Flower, a friendly federated learning framework!\n",
    "\n",
    "In this notebook, we'll simulate a federated learning system with 100 clients. The clients will use TensorFlow/Keras to define model training and evaluation. Let's start by installing Flower (published as `flwr` on PyPI) with the `simulation` extra:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q flwr[\"simulation\"] tensorflow\n",
    "!pip install -q flwr_datasets[\"vision\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's also install Matplotlib so we can make some plots once the simulation is completed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install matplotlib"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we import the required dependencies. The most important imports are Flower (`flwr`) and TensorFlow:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Dict, List, Tuple\n",
    "\n",
    "import tensorflow as tf\n",
    "\n",
    "import flwr as fl\n",
    "from flwr.common import Metrics\n",
    "from flwr.simulation.ray_transport.utils import enable_tf_gpu_growth\n",
    "\n",
    "from datasets import Dataset\n",
    "from flwr_datasets import FederatedDataset\n",
    "\n",
    "VERBOSE = 0\n",
    "NUM_CLIENTS = 100"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's start by defining the model we want to federated. Since we will be working with MNIST, using a fully connected model is sufficient. You can of course customize this model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_model():\n",
    "    \"\"\"Constructs a simple model architecture suitable for MNIST.\"\"\"\n",
    "    model = tf.keras.models.Sequential(\n",
    "        [\n",
    "            tf.keras.layers.Flatten(input_shape=(28, 28)),\n",
    "            tf.keras.layers.Dense(128, activation=\"relu\"),\n",
    "            tf.keras.layers.Dropout(0.2),\n",
    "            tf.keras.layers.Dense(10, activation=\"softmax\"),\n",
    "        ]\n",
    "    )\n",
    "    model.compile(\"adam\", \"sparse_categorical_crossentropy\", metrics=[\"accuracy\"])\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With that out of the way, let's move on to the interesting bits. Federated learning systems consist of a server and multiple clients. In Flower, we create clients by implementing subclasses of `flwr.client.Client` or `flwr.client.NumPyClient`. We use `NumPyClient` in this tutorial because it is easier to implement and requires us to write less boilerplate.\n",
    "\n",
    "To implement the Flower client, we create a subclass of `flwr.client.NumPyClient` and implement the three methods `get_parameters`, `fit`, and `evaluate`:\n",
    "\n",
    "- `get_parameters`: Return the current local model parameters\n",
    "- `fit`: Receive model parameters from the server, train the model parameters on the local data, and return the (updated) model parameters to the server \n",
    "- `evaluate`: Received model parameters from the server, evaluate the model parameters on the local data, and return the evaluation result to the server\n",
    "\n",
    "We mentioned that our clients will use TensorFlow/Keras for the model training and evaluation. Keras models provide methods that make the implementation straightforward: we can update the local model with server-provides parameters through `model.set_weights`, we can train/evaluate the model through `fit/evaluate`, and we can get the updated model parameters through `model.get_weights`.\n",
    "\n",
    "Let's see a simple implementation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FlowerClient(fl.client.NumPyClient):\n",
    "    def __init__(self, trainset, valset) -> None:\n",
    "        # Create model\n",
    "        self.model = get_model()\n",
    "        self.trainset = trainset\n",
    "        self.valset = valset\n",
    "\n",
    "    def get_parameters(self, config):\n",
    "        return self.model.get_weights()\n",
    "\n",
    "    def fit(self, parameters, config):\n",
    "        self.model.set_weights(parameters)\n",
    "        self.model.fit(self.trainset, epochs=1, verbose=VERBOSE)\n",
    "        return self.model.get_weights(), len(self.trainset), {}\n",
    "\n",
    "    def evaluate(self, parameters, config):\n",
    "        self.model.set_weights(parameters)\n",
    "        loss, acc = self.model.evaluate(self.valset, verbose=VERBOSE)\n",
    "        return loss, len(self.valset), {\"accuracy\": acc}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our class `FlowerClient` defines how local training/evaluation will be performed and allows Flower to call the local training/evaluation through `fit` and `evaluate`. Each instance of `FlowerClient` represents a *single client* in our federated learning system. Federated learning systems have multiple clients (otherwise, there's not much to federate, is there?), so each client will be represented by its own instance of `FlowerClient`. If we have, for example, three clients in our workload, we'd have three instances of `FlowerClient`. Flower calls `FlowerClient.fit` on the respective instance when the server selects a particular client for training (and `FlowerClient.evaluate` for evaluation).\n",
    "\n",
    "In this notebook, we want to simulate a federated learning system with 100 clients on a single machine. This means that the server and all 100 clients will live on a single machine and share resources such as CPU, GPU, and memory. Having 100 clients would mean having 100 instances of `FlowerClient` in memory. Doing this on a single machine can quickly exhaust the available memory resources, even if only a subset of these clients participates in a single round of federated learning.\n",
    "\n",
    "In addition to the regular capabilities where server and clients run on multiple machines, Flower, therefore, provides special simulation capabilities that create `FlowerClient` instances only when they are actually necessary for training or evaluation. To enable the Flower framework to create clients when necessary, we need to implement a function called `client_fn` that creates a `FlowerClient` instance on demand. Flower calls `client_fn` whenever it needs an instance of one particular client to call `fit` or `evaluate` (those instances are usually discarded after use). Clients are identified by a client ID, or short `cid`. The `cid` can be used, for example, to load different local data partitions for each client"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now define four auxiliary functions for this example (note the last two are entirely optional):\n",
    "* `get_client_fn()`: Is a function that returns another function. The returned `client_fn` will be executed by Flower's VirtualClientEngine each time a new _virtual_ client (i.e. a client that is simulated in a Python process) needs to be spawn. When are virtual clients spawned? Each time the strategy samples them to do either `fit()` (i.e. train the global model on the local data of a particular client) or `evaluate()` (i.e. evaluate the global model on the validation set of a given client).\n",
    "\n",
    "* `weighted_average()`: This is an optional function to pass to the strategy. It will be executed after an evaluation round (i.e. when client run `evaluate()`) and will aggregate the metrics clients return. In this example, we use this function to compute the weighted average accuracy of clients doing `evaluate()`.\n",
    "\n",
    "* `get_evaluate_fn()`: This is again a function that returns another function. The returned function will be executed by the strategy at the end of a `fit()` round and after a new global model has been obtained after aggregation. This is an optional argument for Flower strategies. In this example, we use the whole MNIST test set to perform this server-side evaluation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_client_fn(dataset: FederatedDataset):\n",
    "    \"\"\"Return a function to construct a client.\n",
    "\n",
    "    The VirtualClientEngine will execute this function whenever a client is sampled by\n",
    "    the strategy to participate.\n",
    "    \"\"\"\n",
    "\n",
    "    def client_fn(cid: str) -> fl.client.Client:\n",
    "        \"\"\"Construct a FlowerClient with its own dataset partition.\"\"\"\n",
    "\n",
    "        # Extract partition for client with id = cid\n",
    "        client_dataset = dataset.load_partition(int(cid), \"train\")\n",
    "\n",
    "        # Now let's split it into train (90%) and validation (10%)\n",
    "        client_dataset_splits = client_dataset.train_test_split(test_size=0.1, seed=42)\n",
    "\n",
    "        trainset = client_dataset_splits[\"train\"].to_tf_dataset(\n",
    "            columns=\"image\", label_cols=\"label\", batch_size=32\n",
    "        )\n",
    "        valset = client_dataset_splits[\"test\"].to_tf_dataset(\n",
    "            columns=\"image\", label_cols=\"label\", batch_size=64\n",
    "        )\n",
    "\n",
    "        # Create and return client\n",
    "        return FlowerClient(trainset, valset).to_client()\n",
    "\n",
    "    return client_fn\n",
    "\n",
    "\n",
    "def weighted_average(metrics: List[Tuple[int, Metrics]]) -> Metrics:\n",
    "    \"\"\"Aggregation function for (federated) evaluation metrics, i.e. those returned by\n",
    "    the client's evaluate() method.\"\"\"\n",
    "    # Multiply accuracy of each client by number of examples used\n",
    "    accuracies = [num_examples * m[\"accuracy\"] for num_examples, m in metrics]\n",
    "    examples = [num_examples for num_examples, _ in metrics]\n",
    "\n",
    "    # Aggregate and return custom metric (weighted average)\n",
    "    return {\"accuracy\": sum(accuracies) / sum(examples)}\n",
    "\n",
    "\n",
    "def get_evaluate_fn(testset: Dataset):\n",
    "    \"\"\"Return an evaluation function for server-side (i.e. centralised) evaluation.\"\"\"\n",
    "\n",
    "    # The `evaluate` function will be called after every round by the strategy\n",
    "    def evaluate(\n",
    "        server_round: int,\n",
    "        parameters: fl.common.NDArrays,\n",
    "        config: Dict[str, fl.common.Scalar],\n",
    "    ):\n",
    "        model = get_model()  # Construct the model\n",
    "        model.set_weights(parameters)  # Update model with the latest parameters\n",
    "        loss, accuracy = model.evaluate(testset, verbose=VERBOSE)\n",
    "        return loss, {\"accuracy\": accuracy}\n",
    "\n",
    "    return evaluate"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now have `FlowerClient` which defines client-side training and evaluation, and `client_fn`, which allows Flower to create `FlowerClient` instances whenever it needs to call `fit` or `evaluate` on one particular client. The last step is to start the actual simulation using `flwr.simulation.start_simulation`. \n",
    "\n",
    "The function `start_simulation` accepts a number of arguments, amongst them the `client_fn` used to create `FlowerClient` instances, the number of clients to simulate `num_clients`, the number of rounds `num_rounds`, and the strategy. The strategy encapsulates the federated learning approach/algorithm, for example, *Federated Averaging* (FedAvg).\n",
    "\n",
    "Flower comes with a number of built-in strategies, but we can also use our own strategy implementations to customize nearly all aspects of the federated learning approach. For this example, we use the built-in `FedAvg` implementation and customize it using a few basic parameters. The last step is the actual call to `start_simulation` which - you guessed it - actually starts the simulation.\n",
    "\n",
    "We can use [Flower Datasets](https://flower.ai/docs/datasets/) to effortlessly obtain an off-the-shelf partitioned dataset or partition one that isn't pre-partitioned. Let's choose MNIST."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Enable GPU growth in your main process\n",
    "enable_tf_gpu_growth()\n",
    "\n",
    "# Download MNIST dataset and partition it\n",
    "mnist_fds = FederatedDataset(dataset=\"mnist\", partitioners={\"train\": NUM_CLIENTS})\n",
    "# Get the whole test set for centralised evaluation\n",
    "centralized_testset = mnist_fds.load_split(\"test\").to_tf_dataset(\n",
    "    columns=\"image\", label_cols=\"label\", batch_size=64\n",
    ")\n",
    "\n",
    "\n",
    "# Create FedAvg strategy\n",
    "strategy = fl.server.strategy.FedAvg(\n",
    "    fraction_fit=0.1,  # Sample 10% of available clients for training\n",
    "    fraction_evaluate=0.05,  # Sample 5% of available clients for evaluation\n",
    "    min_fit_clients=10,  # Never sample less than 10 clients for training\n",
    "    min_evaluate_clients=5,  # Never sample less than 5 clients for evaluation\n",
    "    min_available_clients=int(\n",
    "        NUM_CLIENTS * 0.75\n",
    "    ),  # Wait until at least 75 clients are available\n",
    "    evaluate_metrics_aggregation_fn=weighted_average,  # aggregates federated metrics\n",
    "    evaluate_fn=get_evaluate_fn(centralized_testset),  # global evaluation function\n",
    ")\n",
    "\n",
    "# With a dictionary, you tell Flower's VirtualClientEngine that each\n",
    "# client needs exclusive access to these many resources in order to run\n",
    "client_resources = {\"num_cpus\": 1, \"num_gpus\": 0.0}\n",
    "\n",
    "# Start simulation\n",
    "history = fl.simulation.start_simulation(\n",
    "    client_fn=get_client_fn(mnist_fds),\n",
    "    num_clients=NUM_CLIENTS,\n",
    "    config=fl.server.ServerConfig(num_rounds=10),\n",
    "    strategy=strategy,\n",
    "    client_resources=client_resources,\n",
    "    actor_kwargs={\n",
    "        \"on_actor_init_fn\": enable_tf_gpu_growth  # Enable GPU growth upon actor init.\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can then use the resturned History object to either save the results to disk or do some visualisation (or both of course, or neither if you like chaos). Below you can see how you can plot the centralised accuracy obtainined at the end of each round (including at the very beginning of the experiment) for the global model. This is want the function `evaluate_fn()` that we passed to the strategy reports."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "print(f\"{history.metrics_centralized = }\")\n",
    "\n",
    "global_accuracy_centralised = history.metrics_centralized[\"accuracy\"]\n",
    "round = [data[0] for data in global_accuracy_centralised]\n",
    "acc = [100.0 * data[1] for data in global_accuracy_centralised]\n",
    "plt.plot(round, acc)\n",
    "plt.grid()\n",
    "plt.ylabel(\"Accuracy (%)\")\n",
    "plt.xlabel(\"Round\")\n",
    "plt.title(\"MNIST - IID - 100 clients with 10 clients per round\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Congratulations! With that, you built a Flower client, customized it's instantiation through the `client_fn`, customized the server-side execution through a `FedAvg` strategy configured for this workload, and started a simulation with 100 clients (each holding their own individual partition of the MNIST dataset).\n",
    "\n",
    "Next, you can continue to explore more advanced Flower topics:\n",
    "\n",
    "- Deploy server and clients on different machines using `start_server` and `start_client`\n",
    "- Customize the server-side execution through custom strategies\n",
    "- Customize the client-side execution through `config` dictionaries\n",
    "\n",
    "Get all resources you need!\n",
    "\n",
    "* **[DOCS]** Our complete documenation: https://flower.ai/docs/\n",
    "* **[Examples]** All Flower examples: https://flower.ai/docs/examples/\n",
    "* **[VIDEO]** Our Youtube channel: https://www.youtube.com/@flowerlabs\n",
    "\n",
    "Don't forget to join our Slack channel: https://flower.ai/join-slack/"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "name": "flower.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
